{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "24964024-bf19-46c4-809f-546226e317e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore import context\n",
    "from mindspore.nn import Adam\n",
    "from mindsponge import Sponge\n",
    "from mindsponge import Molecule\n",
    "from mindsponge import ForceField\n",
    "from mindsponge.callback import WriteH5MD, RunInfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "44888a42-9e74-43ae-9744-7c7ba3b0516d",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28b58da8-2137-4a63-ba53-4eddff31cdf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "system = Molecule(template='water.spce.yaml')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "74ec23d9-a8d3-4d91-aa84-bf13859a454f",
   "metadata": {},
   "source": [
    "## 分子系统的复制与扩展\n",
    "\n",
    "在分子对象的最基础的类Molecule中，支持了reduplicate和copy两种复制的模式，其实是同样的操作，不同点在于，reduplicate复制之后会默认把新生成的分子添加到此前定义的分子体系中去。而copy则是创建了一个新的对象，需要通过Molecule内置的append方法，才能把copy出来的对象添加到原始的分子系统中去。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7d5c3022-d3fb-4a37-96a7-4a1158176713",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Molecule<>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "system.reduplicate([0.3, 0, 0])\n",
    "system.reduplicate([0, 0.3, 0])\n",
    "new_sys = system.copy([0, 0, -0.3])\n",
    "system.reduplicate([0, 0, 0.3])\n",
    "system.append(new_sys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f5c47cd-8cf0-4119-b21a-ad005ef50fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential = ForceField(system, parameters='SPCE')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "15883f43-ff0c-4d44-b7d1-27b40378ce88",
   "metadata": {},
   "source": [
    "## 优化器与积分器\n",
    "\n",
    "在MindSPONGE与传统深度学习框架的架构图示意中，我们可以发现分子动力学模拟过程中的积分器与深度学习中的优化器，本质上其实是一回事，都需要获取到gradient然后再去对系统进行更新。而如果在分子动力学模拟的过程中直接使用传统优化器如GradientDescent和Adam等来进行能量最小化，也是一个很重要的过程。在执行动力学模拟过程之前，如果不对系统进行能量最小化，将很有可能在运行过程中出现梯度爆炸或者梯度消失的问题，会为系统的正常运行带来众多的不确定性。\n",
    "\n",
    "![](../../docs/img.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cc5d1c46-19f1-4ef6-866b-1a688e8b5dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = Adam(system.trainable_params(), 1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a65f424a-4733-4a2a-a0fb-93ce6c65e52f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Using fully connected neighbour list (not updated).\n"
     ]
    }
   ],
   "source": [
    "md = Sponge(system, potential, opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5486bed3-ffb6-4391-81a1-7d01b90e5306",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_info = RunInfo(10)\n",
    "cb_h5md = WriteH5MD(system, 'tutorial_b03.h5md', save_freq=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dea67fe3-b2e5-43eb-a9b5-f219b63f2358",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Started simulation at 2023-05-18 15:45:23\n",
      "[MindSPONGE] Step: 0, E_pot: 242.99136\n",
      "[MindSPONGE] Step: 10, E_pot: 163.51727\n",
      "[MindSPONGE] Step: 20, E_pot: 93.06726\n",
      "[MindSPONGE] Step: 30, E_pot: 15.81983\n",
      "[MindSPONGE] Step: 40, E_pot: -66.221664\n",
      "[MindSPONGE] Step: 50, E_pot: -134.5828\n",
      "[MindSPONGE] Step: 60, E_pot: -182.6464\n",
      "[MindSPONGE] Step: 70, E_pot: -220.4002\n",
      "[MindSPONGE] Step: 80, E_pot: -252.21872\n",
      "[MindSPONGE] Step: 90, E_pot: -278.47833\n",
      "[MindSPONGE] Step: 100, E_pot: -302.31818\n",
      "[MindSPONGE] Step: 110, E_pot: -328.98346\n",
      "[MindSPONGE] Step: 120, E_pot: -354.57516\n",
      "[MindSPONGE] Step: 130, E_pot: -373.94397\n",
      "[MindSPONGE] Step: 140, E_pot: -392.10355\n",
      "[MindSPONGE] Step: 150, E_pot: -411.2622\n",
      "[MindSPONGE] Step: 160, E_pot: -430.96765\n",
      "[MindSPONGE] Step: 170, E_pot: -449.49316\n",
      "[MindSPONGE] Step: 180, E_pot: -464.83392\n",
      "[MindSPONGE] Step: 190, E_pot: -476.66644\n",
      "[MindSPONGE] Step: 200, E_pot: -485.10953\n",
      "[MindSPONGE] Step: 210, E_pot: -491.31125\n",
      "[MindSPONGE] Step: 220, E_pot: -496.3717\n",
      "[MindSPONGE] Step: 230, E_pot: -500.64844\n",
      "[MindSPONGE] Step: 240, E_pot: -504.501\n",
      "[MindSPONGE] Step: 250, E_pot: -508.14514\n",
      "[MindSPONGE] Step: 260, E_pot: -511.5323\n",
      "[MindSPONGE] Step: 270, E_pot: -514.3434\n",
      "[MindSPONGE] Step: 280, E_pot: -516.3299\n",
      "[MindSPONGE] Step: 290, E_pot: -517.83\n",
      "[MindSPONGE] Step: 300, E_pot: -519.1944\n",
      "[MindSPONGE] Step: 310, E_pot: -520.51746\n",
      "[MindSPONGE] Step: 320, E_pot: -521.86194\n",
      "[MindSPONGE] Step: 330, E_pot: -523.27686\n",
      "[MindSPONGE] Step: 340, E_pot: -524.8109\n",
      "[MindSPONGE] Step: 350, E_pot: -526.52264\n",
      "[MindSPONGE] Step: 360, E_pot: -528.47986\n",
      "[MindSPONGE] Step: 370, E_pot: -530.74146\n",
      "[MindSPONGE] Step: 380, E_pot: -533.2627\n",
      "[MindSPONGE] Step: 390, E_pot: -535.6827\n",
      "[MindSPONGE] Step: 400, E_pot: -537.2779\n",
      "[MindSPONGE] Step: 410, E_pot: -537.9336\n",
      "[MindSPONGE] Step: 420, E_pot: -538.2575\n",
      "[MindSPONGE] Step: 430, E_pot: -538.4466\n",
      "[MindSPONGE] Step: 440, E_pot: -538.56903\n",
      "[MindSPONGE] Step: 450, E_pot: -538.64294\n",
      "[MindSPONGE] Step: 460, E_pot: -538.6915\n",
      "[MindSPONGE] Step: 470, E_pot: -538.7266\n",
      "[MindSPONGE] Step: 480, E_pot: -538.7523\n",
      "[MindSPONGE] Step: 490, E_pot: -538.7711\n",
      "[MindSPONGE] Step: 500, E_pot: -538.7862\n",
      "[MindSPONGE] Step: 510, E_pot: -538.7986\n",
      "[MindSPONGE] Step: 520, E_pot: -538.80865\n",
      "[MindSPONGE] Step: 530, E_pot: -538.81647\n",
      "[MindSPONGE] Step: 540, E_pot: -538.823\n",
      "[MindSPONGE] Step: 550, E_pot: -538.8281\n",
      "[MindSPONGE] Step: 560, E_pot: -538.8328\n",
      "[MindSPONGE] Step: 570, E_pot: -538.83673\n",
      "[MindSPONGE] Step: 580, E_pot: -538.83954\n",
      "[MindSPONGE] Step: 590, E_pot: -538.8425\n",
      "[MindSPONGE] Step: 600, E_pot: -538.84424\n",
      "[MindSPONGE] Step: 610, E_pot: -538.8356\n",
      "[MindSPONGE] Step: 620, E_pot: -538.8447\n",
      "[MindSPONGE] Step: 630, E_pot: -538.8488\n",
      "[MindSPONGE] Step: 640, E_pot: -538.8504\n",
      "[MindSPONGE] Step: 650, E_pot: -538.8519\n",
      "[MindSPONGE] Step: 660, E_pot: -538.85315\n",
      "[MindSPONGE] Step: 670, E_pot: -538.85425\n",
      "[MindSPONGE] Step: 680, E_pot: -538.85205\n",
      "[MindSPONGE] Step: 690, E_pot: -538.8502\n",
      "[MindSPONGE] Step: 700, E_pot: -538.8557\n",
      "[MindSPONGE] Step: 710, E_pot: -538.85754\n",
      "[MindSPONGE] Step: 720, E_pot: -538.8589\n",
      "[MindSPONGE] Step: 730, E_pot: -538.8602\n",
      "[MindSPONGE] Step: 740, E_pot: -538.86096\n",
      "[MindSPONGE] Step: 750, E_pot: -538.86145\n",
      "[MindSPONGE] Step: 760, E_pot: -538.86224\n",
      "[MindSPONGE] Step: 770, E_pot: -538.8587\n",
      "[MindSPONGE] Step: 780, E_pot: -538.8636\n",
      "[MindSPONGE] Step: 790, E_pot: -538.8617\n",
      "[MindSPONGE] Step: 800, E_pot: -538.8661\n",
      "[MindSPONGE] Step: 810, E_pot: -538.86755\n",
      "[MindSPONGE] Step: 820, E_pot: -538.86914\n",
      "[MindSPONGE] Step: 830, E_pot: -538.86975\n",
      "[MindSPONGE] Step: 840, E_pot: -538.8716\n",
      "[MindSPONGE] Step: 850, E_pot: -538.8723\n",
      "[MindSPONGE] Step: 860, E_pot: -538.861\n",
      "[MindSPONGE] Step: 870, E_pot: -538.8675\n",
      "[MindSPONGE] Step: 880, E_pot: -538.8766\n",
      "[MindSPONGE] Step: 890, E_pot: -538.8781\n",
      "[MindSPONGE] Step: 900, E_pot: -538.8806\n",
      "[MindSPONGE] Step: 910, E_pot: -538.8827\n",
      "[MindSPONGE] Step: 920, E_pot: -538.8766\n",
      "[MindSPONGE] Step: 930, E_pot: -538.88684\n",
      "[MindSPONGE] Step: 940, E_pot: -538.888\n",
      "[MindSPONGE] Step: 950, E_pot: -538.89215\n",
      "[MindSPONGE] Step: 960, E_pot: -538.8959\n",
      "[MindSPONGE] Step: 970, E_pot: -538.89954\n",
      "[MindSPONGE] Step: 980, E_pot: -538.90173\n",
      "[MindSPONGE] Step: 990, E_pot: -538.8978\n",
      "[MindSPONGE] Finished simulation at 2023-05-18 15:45:31\n",
      "[MindSPONGE] Simulation time: 8.51 seconds.\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<mindsponge.core.sponge.Sponge at 0x7f87c692a580>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "md.run(1000, callbacks=[run_info, cb_h5md])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7b1dbf13-3eb4-4cfb-9705-1e7516c37561",
   "metadata": {},
   "source": [
    "![](../../docs/tutorial_b03.gif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore-1.9",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "768f82fbbece3f0acd424cda825de8aa44edbfaa38d8a3df4ff1437f7ae7861f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
